package com.caijb.netty.nio;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.Set;

/**
 * @author caijb
 * @version 1.0
 * @data 2025/8/24 11:04
 */
public class NioServer {
    private final Selector selector;
    private final ServerSocketChannel serverSocketChannel;
    private volatile boolean isRunning = true;

    public NioServer(int port) throws IOException {
        // 创建选择器
        selector = Selector.open();

        // 创建服务器通道
        serverSocketChannel = ServerSocketChannel.open();
        serverSocketChannel.configureBlocking(false); // 设置为非阻塞模式

        // 绑定端口
        serverSocketChannel.bind(new InetSocketAddress(port));

        // 将通道注册到选择器上，监听连接事件
        serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);

        System.out.println("NIO Server started on port: " + port);
    }

    public void start() {
        while (isRunning) {
            try {
                // 等待事件发生，阻塞1秒
                int readyChannels = selector.select();

                if (readyChannels == 0) {
                    continue;
                }

                // 获取有事件发生的SelectionKey集合
                Set<SelectionKey> selectedKeys = selector.selectedKeys();
                Iterator<SelectionKey> keyIterator = selectedKeys.iterator();

                while (keyIterator.hasNext()) {
                    SelectionKey key = keyIterator.next();

                    // 处理完后需要手动移除
                    keyIterator.remove();

                    if (!key.isValid()) {
                        continue;
                    }

                    // 处理不同类型的事件
                    if (key.isAcceptable()) {
                        handleAccept(key);
                    } else if (key.isReadable()) {
                        handleRead(key);
                    } else if (key.isWritable()) {
                        handleWrite(key);
                    }
                }
            } catch (IOException e) {
                System.err.println("Error in server loop: " + e.getMessage());
            }
        }
    }

    // 处理连接事件
    private void handleAccept(SelectionKey key) throws IOException {
        ServerSocketChannel serverChannel = (ServerSocketChannel) key.channel();
        SocketChannel clientChannel = serverChannel.accept();

        if (clientChannel != null) {
            clientChannel.configureBlocking(false);
            // 注册读事件
            clientChannel.register(selector, SelectionKey.OP_READ);
            System.out.println("Client connected: " + clientChannel.getRemoteAddress());
        }
    }

    // 处理读事件
    private void handleRead(SelectionKey key) throws IOException {
        SocketChannel clientChannel = (SocketChannel) key.channel();
        ByteBuffer buffer = ByteBuffer.allocate(1024);

        try {
            int bytesRead = clientChannel.read(buffer);

            if (bytesRead == -1) {
                // 客户端已关闭连接
                System.out.println("Client disconnected: " + clientChannel.getRemoteAddress());
                clientChannel.close();
                key.cancel();
                return;
            }

            if (bytesRead > 0) {
                // 切换为读模式
                buffer.flip();
                byte[] data = new byte[buffer.remaining()];
                buffer.get(data);
                String message = new String(data).trim();

                System.out.println("Received from " + clientChannel.getRemoteAddress() + ": " + message);

                // 准备回显消息
                String response = "Echo: " + message + "\n";
                ByteBuffer responseBuffer = ByteBuffer.wrap(response.getBytes());

                // 注册写事件
                key.attach(responseBuffer);
                key.interestOps(SelectionKey.OP_WRITE);
            }
        } catch (IOException e) {
            System.err.println("Error reading from client: " + e.getMessage());
            clientChannel.close();
            key.cancel();
        }
    }

    // 处理写事件
    private void handleWrite(SelectionKey key) throws IOException {
        SocketChannel clientChannel = (SocketChannel) key.channel();
        ByteBuffer buffer = (ByteBuffer) key.attachment();

        if (buffer != null) {
            clientChannel.write(buffer);

            if (!buffer.hasRemaining()) {
                // 数据已全部写入，重新注册读事件
                key.interestOps(SelectionKey.OP_READ);
            }
        }
    }

    public void stop() throws IOException {
        isRunning = false;
        if (selector != null) {
            selector.close();
        }
        if (serverSocketChannel != null) {
            serverSocketChannel.close();
        }
        System.out.println("NIO Server stopped");
    }

    public static void main(String[] args) {
        try {
            NioServer server = new NioServer(8080);

            // 添加关闭钩子
            Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                try {
                    server.stop();
                } catch (IOException e) {
                    System.err.println("Error stopping server: " + e.getMessage());
                }
            }));

            server.start();
        } catch (IOException e) {
            System.err.println("Failed to start NIO server: " + e.getMessage());
        }
    }
}
